from .metrics import MetricOrValue, get_metric_scalar_value
from .metrics_logger import MetricsLogger


class MLFlowMetricsLogger(MetricsLogger):
    def __init__(self,
        tracking_uri: str | None = None, *,
        experiment_name: str | None = None,
        experiment_id: str | None = None,
        run_name: str = None,
        run_id: str | None = None,
        step_metric_name: str = "step",
    ):
        from mlflow import MlflowClient
        client = MlflowClient(tracking_uri=tracking_uri)

        if run_id is None:
            if experiment_id is None:
                if experiment_name is None:
                    experiment_name = "Default"
                experiment = client.get_experiment_by_name(experiment_name)
                if experiment is None:
                    experiment_id = client.create_experiment(experiment_name)
                else:
                    experiment_id = experiment.experiment_id
            run = client.create_run(experiment_id, run_name=run_name)
            run_id = run.info.run_id

        self.client = client
        self.run_id = run_id
        self.step_metric_name = step_metric_name

    def log(self, metrics: dict[str, MetricOrValue]):
        from mlflow.entities import Metric
        from mlflow.environment_variables import MLFLOW_ENABLE_ASYNC_LOGGING
        from mlflow.utils.time import get_current_time_millis

        metrics_value = {}
        for name, value in metrics.items():
            value = get_metric_scalar_value(value)
            if value is not None:
                metrics_value[name] = value

        step = metrics_value.pop("train/" + self.step_metric_name)

        timestamp = get_current_time_millis()
        synchronous = not MLFLOW_ENABLE_ASYNC_LOGGING.get()
        return self.client.log_batch(
            run_id=self.run_id,
            metrics=[
                Metric(name, value, timestamp, step)
                for name, value in metrics_value.items()],
            synchronous=synchronous)
